#include "TwoTensorFitting.h"

#include <levmar/levmar.h>

#include <teem/ell.h>
#include <teem/ten.h>

#include <stdlib.h>
#include <stdio.h>

#include <math.h>

/* Form a 2D tensor from the parameters */
void PeledRotate2D(float ten[7], float lam1, float lam3, float phi) {
  float cc, ss, d3, d1, d2;

  cc = cos(phi);
  ss = sin(phi);
  d1 = cc*cc*lam1 + ss*ss*lam3;
  d3 = cc*ss*(lam1 - lam3);
  d2 = ss*ss*lam1 + cc*cc*lam3;

  TEN_T_SET(ten, 1.0, d1, d3, 0.0, d2, 0.0, lam3);
  return;
}

/* call back function for levmar optimaization */
void TwoTensorFittingCB(float *pp, float *xx, int mm, int nn, void *_pvlData) {
  float tenA[7], tenB[7];
  int ii;
  AData *pvlData;
  CVector3F *gradients;
 
  AIR_UNUSED(mm);
  pvlData = AIR_CAST(AData *, _pvlData);

  /* Form the tensors using the estimated parms */
  PeledRotate2D(tenA, pp[0], pvlData->lam3, pvlData->phiA);
  PeledRotate2D(tenB, pp[0], pvlData->lam3, pvlData->phiB);

  gradients = AIR_CAST(CVector3F *, pvlData->gradientOrientations);

  for (ii = 0; ii < nn; ii++) {
    float argA, argB, sigA, sigB;
	float egrad[3];
	egrad[0] = gradients[ii].m_x;
	egrad[1] = gradients[ii].m_y;
	egrad[2] = gradients[ii].m_z;
	argA = -pvlData->bValue * TEN_T3V_CONTR(tenA, egrad);
	argB = -pvlData->bValue * TEN_T3V_CONTR(tenB, egrad);
	sigA = exp(argA);
	sigB = exp(argB);

    xx[ii] = pp[1] * sigA + (1.0f - pp[1]) * sigB;
  }
  return;
}

float ConmputeTwoTensorFittingResidual(float &lamda, float *dwi, CVector3F *gradients, const int n, 
	const float bValue, float *eval, float *evec, const float phiA, const float phiB)
{
	float guess[2], loBnd[2], upBnd[2], opts[LM_OPTS_SZ], info[LM_INFO_SZ];
	float tenA[7], tenB[7];

	// set lower and upper bound
	loBnd[0] = 0.0;
	loBnd[1] = 0.0;
	upBnd[0] = eval[0] * 5.0f;
	upBnd[1] = 1.0;

	// set initial guess
	guess[0] = eval[0];
	guess[1] = 0.5f;

	// set constraints for LM optimization
	opts[0] = 1E-3;
	opts[1] = 1E-8;
	opts[2] = 1E-8;
	opts[3] = 1E-8;
	opts[4] = 1E-8;

	AData aData;
	aData.phiA = phiA;
	aData.phiB = phiB;
	aData.bValue = bValue;
	aData.gradientOrientations = gradients;
	aData.lam3 = eval[2];

	int lmret = slevmar_bc_dif(TwoTensorFittingCB, guess, dwi, 2, n, 
		loBnd, upBnd, NULL, 1000, opts, info, NULL, NULL, (void *)&aData);
	
	lamda = guess[0];

	float residual = 0.0f;
	if (-1 == lmret) {
		printf("dlevmar_bc_dif() failed!\n");
	} else {
		/* Form the tensors using the estimated parms */
		PeledRotate2D(tenA, guess[0], eval[2], phiA);
		PeledRotate2D(tenB, guess[0], eval[2], phiB);
		residual = 0.0f;
		for (int i = 0; i < n; i++) {
			float argA, argB, sigA, sigB;
			float egrad[3];
			egrad[0] = gradients[i].m_x;
			egrad[1] = gradients[i].m_y;
			egrad[2] = gradients[i].m_z;
			argA = -bValue * TEN_T3V_CONTR(tenA, egrad);
			argB = -bValue * TEN_T3V_CONTR(tenB, egrad);
			sigA = exp(argA);
			sigB = exp(argB);

			float tempDWI = guess[1] * sigA + (1.0f - guess[1]) * sigB;
			residual += (dwi[i] - tempDWI) * (dwi[i] - tempDWI);
		}
	}
	return residual;
}

CVector3F ComputeDirection(float *eval, float *evec, const float lamda, const float phi)
{
	float ten[7];
	float mat[9], rott[9], matTemp[9];
	// get two tensor fit result
	PeledRotate2D(ten, lamda, eval[2], phi);
	TEN_T2M(mat, ten);

	ELL_3M_TRANSPOSE(rott, evec);
	ELL_3M_MUL(matTemp, mat, evec);
	ELL_3M_MUL(mat, rott, matTemp);

	TEN_M2T(ten, mat);

	double m[9], eigenValues[3], eigenVectors[9], trc, iso[9];
	int ret;
		
	TEN_T2M(m, ten);
	trc = ELL_3M_TRACE(m)/3.0;
	ELL_3M_SCALE_SET(iso, -trc, -trc, -trc);
	ELL_3M_ADD2(m, m, iso);

	ret = ell_3m_eigensolve_d(eigenValues, eigenVectors, m, AIR_TRUE);

	if (ell_cubic_root_single_double == ret) {
		/* this was added to fix a stupid problem with very nearly
		isotropic glyphs, used for demonstration figures */
		if (eigenValues[0] == eigenValues[1]) {
			ELL_3V_CROSS(eigenVectors+6, eigenVectors+0, eigenVectors+3);
		} else {
			ELL_3V_CROSS(eigenVectors+0, eigenVectors+3, eigenVectors+6);
		}
	}

	CVector3F dir;
	dir.m_x = (float)(eigenVectors[0]);
	dir.m_y = (float)(eigenVectors[1]);
	dir.m_z = (float)(eigenVectors[2]);

	return dir;
}